# core/paraformer_engine.py

import logging
import os
import sys
from pathlib import Path
from typing import Dict, Optional, Any, Union
from datetime import datetime
import numpy as np
import torch
import torchaudio

# 尝试导入 FunASR 相关模块
try:
    from funasr import AutoModel
    FUNASR_AVAILABLE = True
except ImportError:
    FUNASR_AVAILABLE = False
    logging.warning("FunASR 未安装，Paraformer 引擎将不可用")

from core.asr_engine import ASREngine
from core.config import settings

logger = logging.getLogger(__name__)


class ParaformerASREngine(ASREngine):
    """
    Paraformer ASR 引擎
    基于 FunASR 框架的中文语音识别模型
    """

    def __init__(self):
        self.model = None
        self._is_loaded = False
        self.model_path = settings.PARAFORMER_MODEL_PATH
        self.device = settings.PARAFORMER_DEVICE
        
        # 检查 FunASR 是否可用
        if not FUNASR_AVAILABLE:
            raise ImportError("FunASR 未安装，无法使用 Paraformer 引擎")
        
        logger.info(f"初始化 Paraformer ASR 引擎 | 模型路径: {self.model_path}")

    def load_model(self) -> None:
        """加载 Paraformer 模型"""
        if self._is_loaded:
            return

        try:
            logger.info(f"正在加载 Paraformer 模型: {self.model_path}")
            
            # 检查模型文件是否存在
            if not self.model_path.exists():
                raise FileNotFoundError(f"Paraformer 模型文件不存在: {self.model_path}")
            
            # 使用 FunASR 的本地模型加载方式
            # 需要指定模型类型和配置文件
            model_dir = str(self.model_path.parent)
            config_file = str(settings.PARAFORMER_CONFIG_PATH)
            
            # 加载模型
            self.model = AutoModel(
                model=model_dir,
                model_revision="v2.0.4",
                device=self.device,
                disable_update=True  # 禁用更新检查
            )
            
            self._is_loaded = True
            logger.info(f"Paraformer 模型加载成功，运行在 {self.device.upper()}")
            
        except Exception as e:
            logger.error(f"Paraformer 模型加载失败: {e}")
            # 尝试使用预训练模型作为备选
            try:
                logger.info("尝试使用预训练 Paraformer 模型...")
                self.model = AutoModel(
                    model="paraformer-zh",
                    model_revision="v2.0.4",
                    device=self.device,
                    disable_update=True
                )
                self._is_loaded = True
                logger.info("预训练 Paraformer 模型加载成功")
            except Exception as e2:
                logger.error(f"预训练模型加载也失败: {e2}")
                raise e

    def transcribe(
        self,
        audio_path: Union[str, Path],
        language: Optional[str] = None,
        task: str = "transcribe",
        return_segments: bool = True
    ) -> Dict[str, Any]:
        """
        执行语音识别
        
        Args:
            audio_path: 音频文件路径
            language: 语言（Paraformer 主要支持中文）
            task: 任务类型（Paraformer 只支持 transcribe）
            return_segments: 是否返回分段信息
            
        Returns:
            包含识别结果的字典
        """
        if not self.is_ready():
            raise RuntimeError("Paraformer 引擎未就绪，请先调用 load_model()")

        # 确保 audio_path 是 Path 对象
        if isinstance(audio_path, str):
            audio_path = Path(audio_path)
            
        if not audio_path.exists():
            raise FileNotFoundError(f"音频文件不存在: {audio_path}")

        try:
            logger.info(f"开始 Paraformer 识别: {audio_path.name}")
            
            # Paraformer 主要支持中文，忽略 language 参数
            if language and language.lower() != "zh":
                logger.warning(f"Paraformer 主要支持中文，忽略语言参数: {language}")
            
            # 执行识别
            result = self.model.generate(input=str(audio_path))
            
            # 解析结果
            if not result or len(result) == 0:
                raise RuntimeError("识别结果为空")
            
            # 提取文本
            text_result = result[0]
            if isinstance(text_result, dict):
                source_text = text_result.get("text", "")
                timestamp = text_result.get("timestamp", [])
            else:
                source_text = str(text_result)
                timestamp = []
            
            # 计算音频时长
            try:
                waveform, sample_rate = torchaudio.load(str(audio_path))
                duration = waveform.shape[1] / sample_rate
            except Exception as e:
                logger.warning(f"无法获取音频时长: {e}")
                duration = 0.0
            
            # 构建结果
            output = {
                "text": source_text,
                "detected_language": "zh",  # Paraformer 主要支持中文
                "language_confidence": 0.95,
                "duration": duration,
                "model": "paraformer-conformer-12e-6d-2048-256-zh",
                "timestamp": datetime.now().isoformat(timespec='seconds'),
                "engine": "paraformer"
            }
            
            # 添加分段信息
            if return_segments and timestamp:
                segments = []
                try:
                    for i, (start, end, text) in enumerate(timestamp):
                        segments.append({
                            "id": i,
                            "start": round(start, 3),
                            "end": round(end, 3),
                            "text": text.strip()
                        })
                    output["segments"] = segments
                except (ValueError, TypeError) as e:
                    logger.warning(f"时间戳解析失败: {e}，使用简单分段")
                    # 如果时间戳格式不正确，使用简单分段
                    output["segments"] = [{
                        "id": 0,
                        "start": 0.0,
                        "end": duration,
                        "text": source_text
                    }]
            elif return_segments:
                # 如果没有时间戳，创建简单的分段
                output["segments"] = [{
                    "id": 0,
                    "start": 0.0,
                    "end": duration,
                    "text": source_text
                }]
            
            logger.info(f"Paraformer 识别完成: {source_text[:60]}...")
            return output
            
        except Exception as e:
            logger.error(f"Paraformer 识别失败 {audio_path}: {e}")
            raise

    def is_ready(self) -> bool:
        """检查引擎是否就绪"""
        return self._is_loaded and self.model is not None

    def get_supported_languages(self) -> Dict[str, list]:
        """获取支持的语言列表"""
        return {
            "source_languages": ["zh"],  # Paraformer 主要支持中文
            "target_languages": [],  # Paraformer 不支持翻译
            "detection_supported": False
        }

    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息"""
        return {
            "name": "Paraformer Conformer",
            "version": "12e-6d-2048-256-zh",
            "framework": "FunASR",
            "language": "zh",
            "device": self.device,
            "model_path": str(self.model_path)
        }


# 检查 FunASR 可用性的辅助函数
def is_funasr_available() -> bool:
    """检查 FunASR 是否可用"""
    return FUNASR_AVAILABLE


def get_funasr_installation_guide() -> str:
    """获取 FunASR 安装指南"""
    return """
    FunASR 安装指南:
    
    1. 安装 FunASR:
       pip install funasr
    
    2. 或者从源码安装:
       git clone https://github.com/alibaba-damo-academy/FunASR.git
       cd FunASR
       pip install -e .
    
    3. 验证安装:
       python -c "from funasr import AutoModel; print('FunASR 安装成功')"
    """
